## fig 3-4
{
  library(Seurat)
  library(SingleCellExperiment)
  library(tidyverse)
  library(ggpubr)
  library(pheatmap)
  options(stringsAsFactors = F)
}

# load data
bgi_sobj <- readRDS("Archive/covid19/data/Final_nCoV_0716_upload_BGI.rds") %>%
  DietSeurat()

blish_sobj <- readRDS("covid19/data/blish_covid_NM2020.rds") %>%
  DietSeurat()

DefaultAssay(bgi_sobj) <- "RNA" # set default assay
DefaultAssay(blish_sobj) <- 'RNA'

# remove doublets --------
mark_doublets <- function(sobj){
  sobj %>%
    as.SingleCellExperiment() %>%
    singleCellTK::runScrublet() %>%
    colData() %>%
    as_tibble(rownames = 'rowname') %>%
    select(c(rowname, contains('scrublet')))
}

join_meta <- function(meta, right){
  meta %>%
    rownames_to_column('rowname') %>%
    left_join(right) %>%
    column_to_rownames('rowname')
}

scr_call <- mark_doublets(blish_sobj)

# save intermediate metadata in case rownames are lost
meta_blish <- join_meta(blish_sobj@meta.data, scr_call)

blish_sobj@meta.data <- meta_blish

scr_call <- mark_doublets(bgi_sobj)

meta_bgi <- join_meta(bgi_sobj@meta.data, scr_call)

nose <- function(x){x[1:5, 1:5]}

nose(meta_bgi)

bgi_sobj@meta.data <- meta_bgi

table(blish_sobj$scrublet_call)
table(bgi_sobj$scrublet_call)

# use singler + monaco to re-annotate cell type
monaco <- celldex::MonacoImmuneData()

blish_sobj %>%
  as.SingleCellExperiment() %>%
  SingleR::SingleR(ref = monaco,
                   labels = monaco$label.fine,
                   clusters = blish_sobj$cell.type) ->
  blish_type_res

tibble(cell.type = blish_type_res@rownames,
       singleR = blish_type_res$pruned.labels) %>%
  mutate(singleR = case_when(
    cell.type %in% c('CD16 Monocyte', 'Platelet', 'RBC') ~ cell.type,
    str_detect(cell.type, 'PB') ~ 'Plasma cells',
    str_detect(cell.type, 'switched') ~ 'Plasma cells',
    TRUE ~ singleR
  ))->
  blishtype

bgi_sobj %>%
  as.SingleCellExperiment() %>%
  SingleR::SingleR(ref = monaco,
                   labels = monaco$label.fine,
                   clusters = bgi_sobj$cell_type) ->
  bgi_type_res

tibble(cell_type = bgi_type_res@rownames,
       singleR = bgi_type_res$pruned.labels) %>%
  mutate(singleR = case_when(
    cell_type == 'Plasma' ~ 'Plasma cells',
    cell_type == 'Monocytes' ~ 'Monocytes',
    TRUE ~ singleR
  )) ->
  bgitype

blish_sobj@meta.data <- join_meta(blish_sobj@meta.data, blishtype)

nose(blish_sobj@meta.data)

bgi_sobj@meta.data <- join_meta(bgi_sobj@meta.data, bgitype)

Idents(bgi_sobj) <- 'singleR'
Idents(blish_sobj) <- 'singleR'

table(Idents(bgi_sobj))

write_rds(blish_sobj, 'DE_cells/data/blish_cleaned.rds')
write_rds(bgi_sobj, 'DE_cells/data/BGI_cleaned.rds')

# T & B cell marker -------------
blish_sobj <- read_rds('DE_cells/data/blish_cleaned.rds')
bgi_sobj <- read_rds('DE_cells/data/BGI_cleaned.rds')

# only interested in healthy human
blish_sobj <- subset(blish_sobj, Status == "Healthy" & scrublet_call == 'Singlet')
bgi_sobj <- subset(bgi_sobj, Stage == "Ctrl" & scrublet_call == 'Singlet')

# soupX to remove ambient RNA from droplet-based data ----
library(SoupX)

toc <- bgi_sobj@assays$RNA@counts

soupProf <- data.frame(
  row.names = rownames(toc),
  est = rowSums(toc)/sum(toc),
  counts = rowSums(toc))

scNoDrops <- toc %>%
  SoupChannel(toc, calcSoupProfile = FALSE) %>%
  setSoupProfile(soupProfile = soupProf) %>%
  setClusters(bgi_sobj$seurat_clusters) %>%
  autoEstCont()

set.seed(42)

out <- adjustCounts(scNoDrops, roundToInt = TRUE)

adj_bgi_sobj <- out %>%
  CreateSeuratObject() %>%
  NormalizeData()

adj_bgi_sobj@meta.data <- bgi_sobj@meta.data %>%
  select(-c(1:3)) %>%
  cbind(adj_bgi_sobj@meta.data)

Idents(adj_bgi_sobj) <- 'singleR'

write_rds(adj_bgi_sobj, 'DE_cells/data/bgi_soupx.rds')

tcr <- c('CD3D','CD3E','CD3G','CD247','TRAC','TRBC1','TRBC2','TRDC','TRGC1','TRGC2')

bcr <- c('CD79A','CD79B','IGKC','IGLC2','IGLC3','IGLC6','IGLC7','IGHA1','IGHA2','IGHG1','IGHG2','IGHG3','IGHG4','IGHD','IGHE','IGHM')

# to avoid a huge tibble, write a 'hook' function
cross_tibble <- expand_grid(
  sobj = c('adj_bgi', 'blish'),
  feature = c('bcr', 'tcr'))

find_hook <- function(hook){
  switch (hook,
          'blish' = blish_sobj,
          'bgi' = bgi_sobj,
          'adj_bgi' = adj_bgi_sobj,
          'bcr' = bcr,
          'tcr' = tcr
  )
}

# violin plot-----------
plot_hooked_violin <- function(sobj,features){
  VlnPlot(find_hook(sobj),
          features = find_hook(features),
          pt.size = 0) &
    theme(axis.title.x = element_blank(),
          text = element_text(size = 5),
          axis.text = element_text(size = 5))}

map2(cross_tibble$sobj,
     cross_tibble$feature,
     plot_hooked_violin) %>%
  lapply(ggplotify::as.ggplot) ->
  violin_list

violin_list[[1]]

# heatmap ------------
plot_hooked_heatmap <- function(sobj, features){
  pos_percent_for_types <- function(type_name){
    find_hook(sobj) %>%
      subset(ident = type_name) %>%
      FetchData(vars = find_hook(features)) %>%
      map(\(x)PercentAbove(x, 0)) %>%
      as.data.frame() %>%
      add_column(cell_type = type_name)
  }
  
  find_hook(sobj) %>%
    Idents() %>%
    unique() %>%
    lapply(pos_percent_for_types) %>%
    purrr::reduce(bind_rows) ->
    heat_mat
  
  heat_mat %>%
    column_to_rownames('cell_type') %>%
    mutate(across(everything(), ~formatC(.*100, format = 'f', digits = 2))) %>%
    mutate(across(everything(), ~paste0(.,'%'))) ->
    percent_mat_chr
  
  heatmat_fig <- heat_mat %>%
    column_to_rownames('cell_type') %>%
    t() %>%
    pheatmap(
      cluster_row = FALSE,
      cluster_cols = FALSE,
      display_numbers = t(percent_mat_chr),
      number_color = "black",
      #legend = FALSE,
      fontsize = 8,
      fontsize_number = case_when(
        str_detect(sobj, 'bgi') ~ 8,
        str_detect(sobj, 'blish') ~ 7
      )
    ) %>%
    ggplotify::as.ggplot()
  
  list(heatmat_fig, percent_mat_chr)}

map2(cross_tibble$sobj,
     cross_tibble$feature,
     plot_hooked_heatmap) %>%
  transpose() ->
  heatmap_list

fig_tibble <- add_column(cross_tibble,
           heatmap = heatmap_list[[1]],
           violin = violin_list)

write_rds(fig_tibble, 'DE_cells/results/fig_tibble.rds')

source_path <- mutate(fig_tibble, path = str_glue('DE_cells/results/fig3-4_heatmap_source_{sobj}_{feature}.csv')) %>%
  pull(path)

walk2(heatmap_list[[2]], source_path, write.csv)

fig_tibble$heatmap[[1]] +
  fig_tibble$violin[[1]] +
  fig_tibble$heatmap[[3]] +
  fig_tibble$violin[[3]] +
  patchwork::plot_annotation(tag_levels = 'A') &
  theme(plot.tag = element_text(size = 32))

ggsave("DE_cells/figures/fig3_ms_bcr_adj.pdf",
       width = 36,
       height = 30,
       units = 'cm')

ggsave("DE_cells/figures/fig3_ms_bcr_adj.png",
       width = 36,
       height = 30,
       units = 'cm')

fig_tibble$heatmap[[2]] +
  fig_tibble$violin[[2]] +
  fig_tibble$heatmap[[4]] +
  fig_tibble$violin[[4]] +
  patchwork::plot_annotation(tag_levels = 'A') &
  theme(plot.tag = element_text(size = 32))

ggsave("DE_cells/figures/fig4_ms_tcr_adj.pdf",
       width = 36,
       height = 30,
       units = 'cm')

ggsave("DE_cells/figures/fig4_ms_tcr_adj.png",
       width = 36,
       height = 30,
       units = 'cm')

